from LeetTool import TreeNode
from LeetTool import build_TreeNode


class Solution:
    def findDistance(self, root: TreeNode, p: int, q: int) -> int:
        ans = 0

        def dfs(node):
            nonlocal ans
            if not node:
                return 0

            left = dfs(node.left)
            right = dfs(node.right)

            if left > 0 and right > 0:
                ans = left + right
                return 0
            elif left > 0 or right > 0:
                if node.val == p or node.val == q:
                    ans = max(left, right)
                    return 0
                else:
                    return max(left, right) + 1
            elif node.val == p and node.val == q:
                ans = 0
                return 0
            elif node.val == p or node.val == q:
                return 1
            else:
                return 0

        dfs(root)
        return ans


if __name__ == "__main__":
    # 3
    print(Solution().findDistance(build_TreeNode([3, 5, 1, 6, 2, 0, 8, None, None, 7, 4]), 5, 0))

    # 2
    print(Solution().findDistance(build_TreeNode([3, 5, 1, 6, 2, 0, 8, None, None, 7, 4]), 5, 7))

    # 0
    print(Solution().findDistance(build_TreeNode([3, 5, 1, 6, 2, 0, 8, None, None, 7, 4]), 5, 5))
